package main

import (
	"flag"
	"fmt"
	"io"
	"net"
	"os"
	"syscall"
	"time"
)

var addr = "127.0.0.1:19771"

func s() {
	p := CreateEpoll()
	defer p.Close()

	conns := make(map[int]*net.TCPConn)

	taddr, err := net.ResolveTCPAddr("tcp", addr)
	if err != nil {
		panic(err)
	}
	l, err := net.ListenTCP("tcp", taddr)
	if err != nil {
		panic(err)
	}

	// Add to epoll
	var lfd int
	rawl, err := l.SyscallConn()
	if err != nil {
		panic(err)
	}
	if err := rawl.Control(func(fd uintptr) {
		lfd = int(fd)
		p.Add(lfd, syscall.EPOLLIN)
	}); err != nil {
		panic(err)
	}

	var buffer [1024]byte
	for {
		events, err := p.Wait(64)
		if err != nil {
			panic(err)
		}

		for _, event := range events {
			switch int(event.Fd) {
			case lfd:
				c, err := l.AcceptTCP()
				switch err {
				case nil:
					rawc, err := c.SyscallConn()
					if err != nil {
						panic(err)
					}

					var cfd int
					if err := rawc.Control(func(fd uintptr) {
						cfd = int(fd)
						syscall.SetNonblock(cfd, true)
					}); err != nil {
						panic(err)
					}

					// Add to epoll
					err = p.Add(cfd, syscall.EPOLLIN)
					if err != nil {
						panic(err)
					}
					conns[cfd] = c

					fmt.Println("Accept Connection")
				case syscall.EAGAIN, syscall.EINTR:
					// retry
				default:
					panic(err)
				}
			default:
				if (event.Events & syscall.EPOLLIN) != 0 {
					c, ok := conns[int(event.Fd)]
					if !ok {
						panic("Unexpected")
					}
					n, err := c.Read(buffer[:])
					switch err {
					case nil:
						fmt.Println("Recv:", string(buffer[:n]))
					case io.EOF:
						fmt.Println("Remote Close Connection")
						delete(conns, int(event.Fd))
						p.Del(int(event.Fd))
						c.Close()
					default:
						panic(err)
					}
				}
			}
		}
	}
}

func c() {
	c, err := net.Dial("tcp", addr)
	if err != nil {
		panic(err)
	}
	defer c.Close()

	for i := 0; i < 3; i++ {
		c.Write([]byte("Emmmmmmmmmm"))
		time.Sleep(time.Second)
	}
}

var (
	sflag = flag.Bool("s", false, "true")
	cflag = flag.Bool("c", false, "false")
)

func init() {
	flag.Parse()
	if !*sflag && !*cflag {
		flag.Usage()
		os.Exit(-1)
	}
}

func main() {
	switch {
	case *sflag:
		s()
	case *cflag:
		c()
	}
}
